import cPickle as pickle
from astro.io import readtxt
from astro.coord import ang_sep
import numpy as np
import matplotlib.pyplot as pl
from matplotlib.mlab import rec_append_fields

import xcorr
import xcorr_CIV
from xcorr_CIV import BINEDGES, NBINS, CBINS
from xcorr import wlya, wlyb, Ckms, cosmo, PC

NRAND = 1000

# maximum separation allowed for an LBG from a QSO sight-line
# (degrees).
MAXSEP = 40. / 60.

# bins to use when measuring mean flux. distances are in Mpc
maxdist = 50.0
#binwidth = 1.0
#radial_binedges = np.arange(0, maxdist + binwidth, binwidth)
#transverse_binedges = np.arange(0, maxdist + binwidth, binwidth)


# linear bins
#BINEDGES = np.arange(0, maxdist + binwidth, binwidth)

# log bins
#BINEDGES = np.logspace(0, 3, 31)

#NBINS = len(BINEDGES) - 1
#NBINS_R = len(radial_binedges) - 1
#NBINS_T = len(transverse_binedges) - 1
#CBINS = BINEDGES[:-1] + BINEDGES[1:] - BINEDGES[:-1]


def readDLAs():
    # read DLA systems
    dla0 = readtxt('/home/nhmc/projects/gas_gal_hiz/dla/dla_all.txt',
                  sep='|', readnames=True)

    const = 1. / (PC * 1.e6) 
    # 1. get the detection ranges for each QSO where DLA can be detected.
    zmin, zmax, los = [], [], []
    for dla in dla0:
        zmax.append( (1 + dla['zem']) * (1 - 3000. / Ckms) - 1 )
        zmin.append( wlyb * (1 + dla['zem']) / wlya - 1 )
        los.append( cosmo.Dc(dla['z']) * const )

    dla1 = rec_append_fields(dla0, ['zmin', 'zmax', 'dlos'],
                             [zmin, zmax, los])

    dla2 = dla1[dla1.z < dla1.zmax]
    return dla2

def make_random_dla(dlas, nrand=None, trypickle=False):
    """ Generate random absorbers. For each QSO containing absorbers,
    generate the number of random absorbers as reals, but with a
    randomly selected redshift

    generate `nrand` times more randoms than reals."""

    random = []
    const = 1. / (PC * 1.e6)

    print 'Generating random DLA lines'
    # this loop takes around 10s to run for nrand=1000
    idnum = 1
    for dla in dlas:
        z0,z1 = dla['zmin'], dla['zmax']
        qsonum = dla['qsonum']
        zran = z0 + (z1 - z0) * np.random.rand(nrand)
        row = [dla[k] for k in 'name ra dec zem'.split()]
        rows = [tuple([i+idnum] + row + [z, cosmo.Dc(z) * const, qsonum])
                for i,z in enumerate(zran)]
        idnum += len(zran)
        random.extend(rows)

    random = np.rec.fromrecords(
        random, names='num,name,ra,dec,zem,z,dlos,qsonum')

    return random

def selectlbgs(lbgs, ra, dec):
    # lbg selection
    c0 = (zmin <= lbgs['z']) & (lbgs['z'] <= zmax)
    lbgs0 = lbgs[c0]
    sep = ang_sep(ra, dec, lbgs0.ra, lbgs0.dec)
    is_close =  sep < MAXSEP
    lbgs1 = lbgs0[is_close]
    # if there are no LBGs close enough, skip to the next QSO.
    if len(lbgs1) == 0:
        print 'No LBGs close enough to %s, skipping' % temp
        return None
    # angular separation of lbgs from QSO sightline (radians)
    lbg_angsep = sep[is_close] * np.pi / 180.
    # comoving transverse separation
    lbg_angdist = np.array([r*los for r,los in
                            zip(lbg_angsep, lbgs1.losdist)])

    return lbgs1, lbg_angsep, lbg_angdist

if 1:
    # read data
    lbgs = xcorr.readlbgs()
    rlbgs = xcorr.read_randomlbgs()
    dla0 = readDLAs()

if 1:
    #CIVa = CIVall[CIVall.isHIRES]
    #CIVa = CIVall[CIVall.rew < 0.3]
    dla1 = dla0
    
    # generate random absorbers
    randoms = make_random_dla(dla1, nrand=NRAND, trypickle=False)

    # loop over every QSO to calculate number of pairs of CIV systems
    # and lbgs as a function of distance
    # need a unique specifier for each qso, filename acts as this
    
    qsos = list(set(dla1.qsonum))

    # ids of CIV absorbers, galaxies that are pair members
#     pairIDs = xcorr_CIV.list_matrix(NBINS_R, NBINS_T)
#     paircounts = np.zeros([NBINS_R, NBINS_T], float)
#     random_pairIDs = xcorr_CIV.list_matrix(NBINS_R, NBINS_T)
#     random_paircounts = np.zeros([NBINS_R, NBINS_T], float)

#     pairIDs_t = [[] for i in range(NBINS_T)]
#     paircounts_t = np.zeros(NBINS_T, float)
#     random_pairIDs_t = [[] for i in range(NBINS_T)]
#     random_paircounts_t = np.zeros(NBINS_T, float)

    pairIDs_s = [[] for i in range(NBINS)]
    paircounts_s = np.zeros(NBINS, float)
    random_pairIDs_s = [[] for i in range(NBINS)]
    random_paircounts_s = np.zeros(NBINS, float)
    # for random lbgs
    random_pairIDs1_s = [[] for i in range(NBINS)]
    random_paircounts1_s = np.zeros(NBINS, float)

    contributing_qsos = []
    for qso in qsos:
        # find all the CIV systems for this QSO
        dla = dla1[dla1.qsonum == qso]
        if len(dla) == 0:
            continue
        zmin, zmax = dla.zmin[0], dla.zmax[0]
        print zmin, zmax
        ra, dec = dla.ra[0], dla.dec[0]

        dlaran = randoms[randoms.qsonum == qso]
        
        # get all the LBGs within
        dv = 5000.
        # go a little further beyond the z limits to look for LBGs.
        zmin = (1 + zmin) * (1 - dv/Ckms) - 1
        zmax = (1 + zmax) * (1 + dv/Ckms) - 1
        print 'Redshift limits for LBGs %5.2f -> %5.2f' % (zmin,zmax)

        temp = selectlbgs(lbgs, ra, dec)
        if temp is None:
            continue
        lbgs1, lbg_angsep, lbg_angdist = temp 
        contributing_qsos.append( (ra,dec,qso) )

        temp = selectlbgs(rlbgs, ra, dec)
        if temp is None:
            continue
        rlbgs1, rlbg_angsep, rlbg_angdist = temp 

        ######################################
        # for xi(r,t)
        ######################################
        
#         # real pairs
#         pairs, pairids = xcorr_CIV.calc_npairs(
#             radial_binedges, transverse_binedges, lbg_angdist,
#             lbgs1.losdist, lbgs1.num, dla.dlos, dla.num)
#         paircounts += pairs
#         xcorr_CIV.list_matrix_extend(pairIDs, pairids)

#         # now random
#         pairs, pairids = xcorr_CIV.calc_npairs(
#             radial_binedges, transverse_binedges, lbg_angdist,
#             lbgs1.losdist, lbgs1.num, dlaran.dlos, dlaran.num)
#         random_paircounts += pairs
#         xcorr_CIV.list_matrix_extend(random_pairIDs, pairids)

        ######################################
        # for xi(s)
        ######################################
        pairs, pairids = xcorr_CIV.calc_npairs_s(
            BINEDGES, lbg_angdist,
            lbgs1.losdist, lbgs1.num, dla.dlos, dla.dlanum,
            maxdist)
        paircounts_s += pairs
        for j,p in enumerate(pairids):
            pairIDs_s[j].extend(p)
         
        pairs, pairids = xcorr_CIV.calc_npairs_s(
            BINEDGES, lbg_angdist,
            lbgs1.losdist, lbgs1.num, dlaran.dlos, dlaran.num,
            maxdist)
        random_paircounts_s += pairs
        for j,p in enumerate(pairids):
            random_pairIDs_s[j].extend(p)

        pairs, pairids = xcorr_CIV.calc_npairs_s(
            BINEDGES, rlbg_angdist,
            rlbgs1.losdist, rlbgs1.num, dla.dlos, dla.dlanum,
            maxdist)
        random_paircounts1_s += pairs
        for j,p in enumerate(pairids):
            random_pairIDs1_s[j].extend(p)
        ######################################
        # for xi(t)
        ######################################
#         pairs, pairids = xcorr_CIV.calc_npairs_t(
#             transverse_binedges, lbg_angdist,
#             lbgs1.losdist, lbgs1.num, dla.dlos, dla.num)
#         paircounts_t += pairs
#         xcorr_CIV.list_matrix_extend(pairIDs_t, pairids)
        
#         pairs, pairids = xcorr_CIV.calc_npairs_t(
#             transverse_binedges, lbg_angdist,
#             lbgs1.losdist, lbgs1.num, dlaran.dlos, dlaran.num)
#         random_paircounts_t += pairs
#         xcorr_CIV.list_matrix_extend(random_pairIDs_t, pairids)

if 1:
    #print contributing_qsos

    # look at xi(s) first...
    xcorr_CIV.plot1d(pairIDs_s, random_pairIDs1_s, cbins=CBINS,
                     title=r'DLA-LBG $\xi(s)$ random lbgs', nrandom=20,
                     prefix='DLAlbg_xi_s_logbins_rlbg')

    idgals = xcorr.unique_lbgids(pairIDs_s)
    c4gals = xcorr.unique_absids(pairIDs_s)


if 0:
    # now find xi
    xi = NRAND * paircounts / random_paircounts - 1

    xcorr_CIV.plotmatrix(paircounts, prefix='CIVpaircounts',
               title='Number of LBG - real CIV pairs')
    xcorr_CIV.plotmatrix(random_paircounts / float(NRAND),
               prefix='CIVrandom_paircounts',
               title='Average number of LBG - random CIV pairs')
    xcorr_CIV.plotmatrix(np.nan_to_num(xi), vmin=-1, prefix='CIVxi',
               title=r'$\xi(r_{\theta}$, $r_{z})$ (real / random - 1)')
    
    xi_t = NRAND * paircounts_t / random_paircounts_t - 1

    xcorr_CIV.plot1d(pairIDs_s, random_pairIDs_s, title=r'CIV-LBG $\xi(s)$',
                     prefix='CIVlbg_xi_s')
    show()


